#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <getopt.h>

#include <unistd.h>
#include <errno.h>
#include <signal.h>
#include <fcntl.h>
#include <sys/time.h>

#include <event2/event.h>
#include <event2/event_struct.h>

#if (HAVE_CONFIG_H)
#include "config.h"
#endif

#include "zip_utils.h"
#include "zip_dns_lib.h"
#include "zip_dns_server.h"
#include "dig_mem_pool.h"
#include "dig_wire_name.h"
#include "zip_record_store.h"
#include "zip_pkg_manager.h"
#include "zip_pkg_list.h"
#include "dig_thread_pool.h"
#include "pthread.h"

#define INIT_MEM_SIZE 1000000

#define MAX_QUERY_BUF_LEN 512


#ifndef PATH_MAX
#define PATH_MAX 1024
#endif
static pthread_rwlock_t rw_lock;
typedef struct dnscache
{
    struct event_base *base_event_;
    struct event *sig_int_event_;
    struct event *sig_term_event_;
    dns_server_t *dns_server_;
    record_store_t *record_store_;
    mem_pool_t *mp_;
    thread_pool_t *tp_;
}dnscache_t;

void 
dnscache_handle_response(const char *raw_data , uint32_t data_len, void *arg);

static void 
dnscache_delete(dnscache_t *dnscache);

static void
dnscache_init_sig(dnscache_t *dnscache);

static void *
dnscache_handle_query_pkg(void *arg);

static void *
dnscache_handle_response_pkg(void *arg);

static void
stop_call_back(evutil_socket_t fd, short event, void *arg)
{
    dnscache_t *dnscache = (dnscache_t*)arg;
    event_del(dnscache->sig_int_event_);
    event_del(dnscache->sig_term_event_);
    dns_server_stop(dnscache->dns_server_);
    thread_pool_stop(dnscache->tp_);
}

typedef struct query_session
{
    uint8_t raw_data[UDP_MAX_PACKAGE_LEN];
    uint16_t data_len;
    socket_t client_socket;
    addr_t client_addr;
    wire_name_t *wire_name;
    dnscache_t *dnscache;
}query_session_t;

static void
handle_dns_query(const uint8_t *raw_data,
                 uint16_t data_len,
                 socket_t *client_socket,
                 const addr_t *client_addr,
                 void *arg)
{

    if (data_len < 12 + 4 + 1 || data_len > 512)
        return ;

    if (dns_header_get_opt(raw_data) != QUERY_STAND)
        return ;

    dnscache_t *dnscache = (dnscache_t *)arg;
    query_session_t *session = mem_pool_alloc(dnscache->mp_);
    if (!session)
       return;

    memcpy(session->raw_data, raw_data, data_len);
    session->data_len = data_len;
    session->client_socket = *client_socket;
    session->client_addr = *client_addr;
    session->dnscache = dnscache;
    buffer_t buf;
    buffer_create_from(&buf, (void*)(raw_data + 12), data_len - 12);
    session->wire_name = wire_name_from_wire(&buf);
    if (session->wire_name)
    {
        if (DNS_QUERY == pkg_query_or_response(raw_data))
        {
            dnscache_handle_query_pkg(session);
            //thread_pool_add_task_to_run(dnscache->tp_, dnscache_handle_query_pkg, session);
        }
        else
        {
            dnscache_handle_response_pkg(session);
           // thread_pool_add_task_to_run(dnscache->tp_, dnscache_handle_response_pkg, session);
        }
    }
}

static void *
dnscache_handle_query_pkg(void *arg)
{
    query_session_t *session = (query_session_t *)arg;
    uint16_t query_type = pkg_get_query_type(session->raw_data, session->data_len);
    uint16_t id = pkg_get_query_id(session->raw_data);
    pkg_data_t *pkg = record_store_get_pkg(session->dnscache->record_store_,
                                           session->wire_name,
                                           query_type);
    if (pkg)
    {
        session->data_len = pkg_data_get_len(pkg);
        memcpy(session->raw_data, pkg_data_get_data(pkg), session->data_len);
        pkg_set_response_id(session->raw_data, id);

        socket_write_to(&session->client_socket,
                        (uint8_t *)session->raw_data,
                        session->data_len,
                        &session->client_addr);

        mem_pool_free(session->dnscache->mp_, session);
        wire_name_delete(session->wire_name);

        return;
    }

    uint8_t *header = (uint8_t *)session->raw_data;
    uint16_t flag = *((uint16_t *)header + 1);
    *((uint16_t *)header + 1) = flag | 128;

    socket_write_to(&session->client_socket,
                    (uint8_t *)session->raw_data,
                    session->data_len,
                    &session->client_addr);

    wire_name_delete(session->wire_name);
    mem_pool_free(session->dnscache->mp_, session);
    return;
}

static void *
dnscache_handle_response_pkg(void *arg)
{
    query_session_t *session = (query_session_t *)arg;

    uint16_t query_type = pkg_get_query_type(session->raw_data, session->data_len);
    pkg_data_t *pkg = pkg_data_create_from(session->raw_data, session->data_len, query_type);
    if (pkg)
    {
        if (record_store_get_total_mem_size(session->dnscache->record_store_) > 2000)
            record_store_free_part_momery(session->dnscache->record_store_, 2000);
        record_store_insert_pkg(session->dnscache->record_store_, session->wire_name, pkg);
    }
    wire_name_delete(session->wire_name);
    mem_pool_free(session->dnscache->mp_, session);
}

static void
dns_cache_print_help()
{
    fprintf(stderr, "Usage: dns_cache [OPTION]...\n");
    fprintf(stderr,
            "Supported options:\n"
            "  -s             bind addr.\n"
            "  -p             listen on port.\n"
            "  -n             startup thread numbers.\n"
            "  -v             Print version information.\n");

    fprintf(stderr, "Version %s. Report bugs to <%s>.\n",
            PACKAGE_VERSION, PACKAGE_BUGREPORT);

}

static void
dns_cache_print_version(void)
{
    fprintf(stderr, "%s version %s\n", PACKAGE_NAME, PACKAGE_VERSION);
    fprintf(stderr, "Written by knet.\n");
    fprintf(stderr,
            "Copyright (C) 2011 knet.\n");
    exit(0);
}

dnscache_t * 
dnscache_create(char *server_addr, uint16_t port)
{
    dnscache_t *dnscache = malloc(sizeof(dnscache_t));
    if (!dnscache)
        return NULL;
    dnscache->base_event_ = event_base_new();

    dnscache_init_sig(dnscache);

    dnscache->mp_ = mem_pool_create(sizeof(query_session_t), 1000);
    ASSERT(dnscache->mp_, "mem pool create failed");

    dnscache->tp_ = thread_pool_create(1);
    ASSERT(dnscache->tp_, "thread pool create failed");


    dnscache->record_store_ = record_store_create();
    ASSERT(dnscache->record_store_, "create record store failed");


    dnscache->dns_server_ = dns_server_create(dnscache->base_event_, 
                                     UDP_SERVER, 
                                     server_addr, 
                                     port);
    dns_server_start(dnscache->dns_server_, handle_dns_query, dnscache);

    pthread_rwlock_init(&rw_lock, NULL);

    return dnscache;
}

static void
dnscache_run(dnscache_t *dnscache)
{
    thread_pool_start(dnscache->tp_);
    event_base_dispatch(dnscache->base_event_);
}

static void
dnscache_delete(dnscache_t *dnscache)
{
    event_free(dnscache->sig_int_event_);
    event_free(dnscache->sig_term_event_);
    dns_server_delete(dnscache->dns_server_);
    record_store_delete(dnscache->record_store_);
    mem_pool_delete(dnscache->mp_);
    thread_pool_delete(dnscache->tp_);
    pthread_rwlock_destroy(&rw_lock);
    free(dnscache);
}

static void
dnscache_init_sig(dnscache_t *dnscache)
{
    struct sigaction sighandler;
    sighandler.sa_handler = SIG_IGN;
    sighandler.sa_flags = 0;
    sigfillset(&sighandler.sa_mask);

    int ret = sigaction(SIGPIPE, &sighandler, NULL);
    ASSERT (ret == 0, "sigaction failed");
    ret = sigaction(SIGHUP, &sighandler, NULL);
    ASSERT (ret == 0, "sigaction failed");

    dnscache->sig_int_event_ = evsignal_new(dnscache->base_event_, 
                                  SIGINT, 
                                  stop_call_back, 
                                  dnscache);

    dnscache->sig_term_event_ = evsignal_new(dnscache->base_event_, 
                                    SIGTERM, 
                                    stop_call_back, 
                                    dnscache);
    event_add(dnscache->sig_int_event_, NULL);
    event_add(dnscache->sig_term_event_, NULL);
}

int main(int argc, char *argv[])
{
    int thread_num = 1;

    char server_addr[MAX_IP_STR_LEN] = "0.0.0.0";
    uint16_t port = 6053;

    int opt = 0;
    while ((opt = getopt(argc, argv, "s:p:n:")) != -1)
    {
        switch(opt)
        {
            case 's':
                strncpy(server_addr, optarg, MAX_IP_STR_LEN);
                break;
            case 'p':
                port = strtol(optarg, NULL, 10);
                break;
            case 'h':
                dns_cache_print_help();
                exit(0);
            case 'n':
                thread_num = strtol(optarg, NULL, 10);
                break;
            case 'v':
                dns_cache_print_version();
                exit(0);
            default:
                dns_cache_print_help();
                exit(0);
        }
    }

    dnscache_t *dnscache = dnscache_create(server_addr, port);
    dnscache_run(dnscache);
    dnscache_delete(dnscache);
    return 0;
}

